# -*- coding: utf-8 -*-

import os
import pandas as pd
from sklearn.model_selection import StratifiedKFold
from keras.utils import np_utils
from keras.callbacks import EarlyStopping, ModelCheckpoint

from lstm_model import make_tokenizer, make_embedding_matrix, pad_sequence
from leaky_lstm_model import cnn_lstm_leak_model


if __name__ == '__main__':
    appMaxLen = 174  # 174
    appEmbedDim = 50
    appTknz = make_tokenizer(os.path.join("data", "deviceid_packages_cooked.txt"),
                             os.path.join("cache", "app.tokenizer"))
    appEmbeddings = make_embedding_matrix(appTknz, appEmbedDim, os.path.join("cache", "app_cbow_50d.vec"),
                                          os.path.join("cache", "app"))

    brandMaxLen = 10
    brandEmbedDim = 5
    brandTknz = make_tokenizer(os.path.join("data", "deviceid_brand_cooked.txt"),
                               os.path.join("cache", "brand.tokenizer"))
    brandEmbeddings = make_embedding_matrix(brandTknz, brandEmbedDim, os.path.join("cache", "brand_cbow_5d.vec"),
                                            os.path.join("cache", "brand"))

    names = ["app", "brand", "label", "gender", "age"]
    df = pd.read_csv(os.path.join("data", "app_brand.train"), sep=",", names=names)
    print df.shape

    leak_df = pd.read_csv(os.path.join("data", "app_gender.train"), sep=",")
    # leak_df.mean().to_pickle(os.path.join("cache", "age_mean.pkl"))
    # leak_df.std().to_pickle(os.path.join("cache", "age_std.pkl"))
    leak_df = (leak_df - leak_df.mean()) / leak_df.std()
    leak_names = leak_df.columns.values
    leakDim = leak_df.shape[1]
    print leak_df.shape

    df = pd.concat([df, leak_df], axis=1)
    print df.shape

    label = df.pop("label")
    # y = np_utils.to_categorical(label)

    kfold = StratifiedKFold(n_splits=5, random_state=42, shuffle=True)
    for i, (train_index, valid_index) in enumerate(kfold.split(df, label)):
        print "[fold-{}]".format(i)
        train_X, valid_X = df.iloc[train_index], df.iloc[valid_index]
        train_y, valid_y = np_utils.to_categorical(label[train_index]), np_utils.to_categorical(label[valid_index])

        train_input = [pad_sequence(train_X["app"], appTknz, appMaxLen),
                       pad_sequence(train_X["brand"], brandTknz, brandMaxLen),
                       train_X[leak_names]]
        valid_input = [pad_sequence(valid_X["app"], appTknz, appMaxLen),
                       pad_sequence(valid_X["brand"], brandTknz, brandMaxLen),
                       valid_X[leak_names]]
        cnnLstm = cnn_lstm_leak_model(
            appTknz, appEmbedDim, appEmbeddings, appMaxLen,
            brandTknz, brandEmbedDim, brandEmbeddings, brandMaxLen,
            leakDim,
            num_class=22
        )

        earlyStop = EarlyStopping(monitor="val_loss", patience=3, verbose=0)
        checkPoint = ModelCheckpoint("model/cnn_lstm_leak_{}.h5".format(i), monitor="val_loss", save_best_only=True, verbose=0)

        cnnLstm.fit(train_input, train_y, validation_data=(valid_input, valid_y), epochs=15, batch_size=64,
                    callbacks=[earlyStop, checkPoint])
